/*
 * @Author: liusheng
 * @Date: 2022-05-10 22:19:24
 * @LastEditors: liusheng
 * @LastEditTime: 2022-05-14 20:24:49
 * @Description: 剑指 Offer II 061. 和最小的 k 个数对
 * email:liusheng613@126.com
 * Copyright (c) 2022 by liusheng/liusheng, All Rights Reserved. 
 * 
 * 剑指 Offer II 061. 和最小的 k 个数对
给定两个以升序排列的整数数组 nums1 和 nums2 , 以及一个整数 k 。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。

请找到和最小的 k 个数对 (u1,v1),  (u2,v2)  ...  (uk,vk) 。

 

示例 1:

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数：
    [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
示例 2:

输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数：
     [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
示例 3:

输入: nums1 = [1,2], nums2 = [3], k = 3 
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]
 

提示:

1 <= nums1.length, nums2.length <= 104
-109 <= nums1[i], nums2[i] <= 109
nums1, nums2 均为升序排列
1 <= k <= 1000
 

注意：本题与主站 373 题相同：https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/
 */

#include "header.h"

class Solution {
public:
    /* using priority queue to keep the topk small sums
    using visited to record wether a pair is visited
    priority queue only push the pair of index from nums1,nums2
    every iterator get the top queue value(smallest index pair),then pop it,add two adjacent pair into the queue
    */
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        int m = nums1.size();
        int n = nums2.size();
        vector<vector<int>> topKPair;
        vector<vector<bool>> visited(m,vector<bool>(n,false));
        using pd = pair<int,int>;
        auto pairGreaterCmp = [&] (const pd & a,const pd & b)
        {
            return (nums1[a.first] + nums2[a.second]) > (nums1[b.first] + nums2[b.second]);
        };

        priority_queue<pair<int,int>,vector<pair<int,int>>,decltype(pairGreaterCmp)> littleEndienQ(pairGreaterCmp);
        littleEndienQ.push(make_pair(0,0));
        while (topKPair.size() < k && !littleEndienQ.empty())
        {
            auto tpV = littleEndienQ.top();
            littleEndienQ.pop();
            int p1 = tpV.first;
            int p2 = tpV.second;
            topKPair.push_back({nums1[p1],nums2[p2]});
            if (p1 + 1 < m && !visited[p1+1][p2])
            {
                littleEndienQ.push(make_pair(p1+1,p2));
                visited[p1+1][p2] = true;
            }
            if (p2 + 1 < n && !visited[p1][p2 + 1])
            {
                littleEndienQ.push(make_pair(p1,p2+1));
                visited[p1][p2+1] = true;
            }
        }

        return topKPair;
    }
};

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        int m = nums1.size();
        int n = nums2.size();
        
        using pd = pair<int,int>;
        auto pairGreaterCmp = [&] (const pd & a,const pd & b) 
        {
            return (nums1[a.first] + nums2[a.second]) > (nums1[b.first] + nums2[b.second]);
        };
        
        /*this won't compile*/
        // auto pairHash = [] (const pd & pr)
        // {
        //     return pr.first ^ pr.second;
        // };
        
        priority_queue<pd,vector<pd>,decltype(pairGreaterCmp)> littleEndianPairQ(pairGreaterCmp);
        struct pairHash
        {
            size_t operator() (const pd & pr) const
            {
                return pr.first ^ pr.second;
            } 
        };
        //visited marked the pair has been dealed(push into priority queue)
        unordered_set<pd,pairHash> visited;
        //insert the minist pair,there must be (0,0)
        littleEndianPairQ.push(make_pair(0,0));
      
        visited.insert(make_pair(0,0));
        vector<vector<int>> topkSmallPair;
        while (topkSmallPair.size() < k && !littleEndianPairQ.empty())
        {
            pd smallIndexPair = littleEndianPairQ.top();
            littleEndianPairQ.pop();
            
            int p1 = smallIndexPair.first;
            int p2 = smallIndexPair.second;
            topkSmallPair.push_back({nums1[p1],nums2[p2]});
            
            //push the two adjacent pair into queue
            if (!visited.count(make_pair(p1 + 1,p2)) && p1 + 1 < m)
            {
                littleEndianPairQ.push(make_pair(p1 + 1,p2));
                visited.insert(make_pair(p1 + 1,p2));
            }
            
            if (!visited.count(make_pair(p1,p2 + 1)) && p2 + 1 < n )
            {
                littleEndianPairQ.push(make_pair(p1,p2 + 1));
                visited.insert(make_pair(p1,p2 + 1));
            }
        }
        
        return topkSmallPair;
    }
};

/*
add nums1's k pair (0,0),(1,0) .... (k-1,0) to priority queue
then every time retrive element from queue,we only add the nums2's index
this can avoid added element repeatlly
*/
class Solution {
public:
    /* using priority queue to keep the topk small sums
    using visited to record wether a pair is visited
    priority queue only push the pair of index from nums1,nums2
    every iterator get the top queue value(smallest index pair),then pop it,add two adjacent pair into the queue
    */
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        int m = nums1.size();
        int n = nums2.size();
        using PairIndex = pair<int,int>;

        auto pairGreaterCmp = [&] (const PairIndex & a,const PairIndex & b)
        {
            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
        };

        priority_queue<PairIndex,vector<PairIndex>,decltype(pairGreaterCmp)> topKIndexPair(pairGreaterCmp);
        vector<vector<int>> topKIndex;
        /* add most minist of nums1.size() and k ,so when pop out,
        just add the nums2's index
        */
        int minInitialNum = min(nums1.size(),(size_t)k);
        for (int i = 0; i < minInitialNum; ++i)
        {
            topKIndexPair.push(make_pair(i,0));
        }

        while (topKIndex.size() < k && !topKIndexPair.empty())
        {
            int index1 = topKIndexPair.top().first;
            int index2 = topKIndexPair.top().second;
            topKIndexPair.pop();

            topKIndex.push_back({nums1[index1],nums2[index2]});

            //since index1 has been all added to priorityQ,only add nums2's index
            if (index2 + 1 < n)
            {
                topKIndexPair.push(make_pair(index1,index2 + 1));
            }
        }

        return topKIndex;
    }
};
